import pygame
import random
import sys
import math
from enum import Enum

# 初始化pygame
pygame.init()

# 游戏设置
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
PUZZLE_SIZE = 400  # 拼图区域大小
GRID_SIZE = 3  # 3x3拼图
PIECE_SIZE = PUZZLE_SIZE // GRID_SIZE
MARGIN = 50  # 边距
FPS = 60

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (128, 128, 128)
LIGHT_GRAY = (200, 200, 200)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)
DARK_GRAY = (64, 64, 64)

# 游戏状态
class GameState(Enum):
    MENU = 1
    PLAYING = 2
    WON = 3
    SELECT_IMAGE = 4

class PuzzlePiece:
    def __init__(self, image, correct_pos, current_pos, piece_id):
        self.image = image
        self.correct_pos = correct_pos  # (row, col)
        self.current_pos = current_pos  # (row, col)
        self.piece_id = piece_id
        self.rect = pygame.Rect(
            MARGIN + current_pos[1] * PIECE_SIZE,
            MARGIN + current_pos[0] * PIECE_SIZE,
            PIECE_SIZE,
            PIECE_SIZE
        )
        self.is_dragging = False
        self.drag_offset = (0, 0)
        self.is_correct = (correct_pos == current_pos)

    def update_position(self, new_pos):
        self.current_pos = new_pos
        self.rect = pygame.Rect(
            MARGIN + new_pos[1] * PIECE_SIZE,
            MARGIN + new_pos[0] * PIECE_SIZE,
            PIECE_SIZE,
            PIECE_SIZE
        )
        self.is_correct = (self.correct_pos == new_pos)

    def draw(self, screen):
        if self.is_dragging:
            screen.blit(self.image, self.rect)
        else:
            screen.blit(self.image, self.rect)

        # 绘制边框
        border_color = GREEN if self.is_correct else BLACK
        pygame.draw.rect(screen, border_color, self.rect, 2)

    def draw_at_position(self, screen, pos):
        rect = pygame.Rect(pos[0], pos[1], PIECE_SIZE, PIECE_SIZE)
        screen.blit(self.image, rect)
        pygame.draw.rect(screen, BLACK, rect, 2)

class PuzzleGame:
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("拼图游戏 - Puzzle Game")
        self.clock = pygame.time.Clock()

        # 字体
        self.font_large = pygame.font.Font(None, 48)
        self.font_medium = pygame.font.Font(None, 36)
        self.font_small = pygame.font.Font(None, 24)

        # 游戏状态
        self.state = GameState.MENU
        self.pieces = []
        self.moves = 0
        self.time_elapsed = 0
        self.game_started = False

        # 拖拽状态
        self.dragging_piece = None
        self.drag_start_pos = None

        # 预设图片（用彩色方块代替实际图片）
        self.create_default_image()

    def create_default_image(self):
        """创建默认的彩色渐变图片"""
        self.original_image = pygame.Surface((PUZZLE_SIZE, PUZZLE_SIZE))

        # 创建彩色渐变图案
        for row in range(GRID_SIZE):
            for col in range(GRID_SIZE):
                # 为每个小块创建不同的颜色
                color = (
                    (row * 255 // GRID_SIZE + col * 50) % 256,
                    (col * 255 // GRID_SIZE + row * 50) % 256,
                    ((row + col) * 255 // (GRID_SIZE * 2)) % 256
                )

                rect = pygame.Rect(
                    col * PIECE_SIZE,
                    row * PIECE_SIZE,
                    PIECE_SIZE,
                    PIECE_SIZE
                )
                pygame.draw.rect(self.original_image, color, rect)

                # 在每个小块上绘制数字
                number = row * GRID_SIZE + col + 1
                text = self.font_large.render(str(number), True, WHITE)
                text_rect = text.get_rect(center=rect.center)
                self.original_image.blit(text, text_rect)

                # 绘制图案装饰
                if row == col:  # 对角线
                    pygame.draw.circle(self.original_image, WHITE, rect.center, 20, 3)
                if row + col == GRID_SIZE - 1:  # 反对角线
                    pygame.draw.rect(self.original_image, WHITE, rect.inflate(-40, -40), 3)

    def create_puzzle_pieces(self):
        """创建拼图块"""
        self.pieces = []
        piece_id = 0

        for row in range(GRID_SIZE):
            for col in range(GRID_SIZE):
                # 从原图中切割小块
                piece_rect = pygame.Rect(
                    col * PIECE_SIZE,
                    row * PIECE_SIZE,
                    PIECE_SIZE,
                    PIECE_SIZE
                )
                piece_image = self.original_image.subsurface(piece_rect).copy()

                piece = PuzzlePiece(
                    piece_image,
                    (row, col),  # 正确位置
                    (row, col),  # 当前位置
                    piece_id
                )
                self.pieces.append(piece)
                piece_id += 1

    def shuffle_pieces(self):
        """打乱拼图块"""
        positions = [(row, col) for row in range(GRID_SIZE) for col in range(GRID_SIZE)]
        random.shuffle(positions)

        for i, piece in enumerate(self.pieces):
            piece.update_position(positions[i])

    def get_piece_at_pos(self, pos):
        """获取指定位置的拼图块"""
        for piece in self.pieces:
            if piece.rect.collidepoint(pos) and not piece.is_dragging:
                return piece
        return None

    def get_empty_position(self, exclude_pos):
        """获取空位置（用于交换）"""
        all_positions = {(row, col) for row in range(GRID_SIZE) for col in range(GRID_SIZE)}
        occupied_positions = {piece.current_pos for piece in self.pieces if piece.current_pos != exclude_pos}
        empty_positions = all_positions - occupied_positions
        return empty_positions.pop() if empty_positions else None

    def swap_pieces(self, piece1, piece2):
        """交换两个拼图块的位置"""
        pos1 = piece1.current_pos
        pos2 = piece2.current_pos
        piece1.update_position(pos2)
        piece2.update_position(pos1)
        self.moves += 1

    def check_win(self):
        """检查是否完成拼图"""
        return all(piece.is_correct for piece in self.pieces)

    def draw_menu(self):
        """绘制主菜单"""
        self.screen.fill(WHITE)

        # 标题
        title = self.font_large.render("PUZZLE GAME", True, BLACK)
        title_rect = title.get_rect(center=(SCREEN_WIDTH // 2, 150))
        self.screen.blit(title, title_rect)

        # 说明
        instructions = [
            "Drag and drop puzzle pieces",
            "to arrange them in correct order",
            "",
            "Press SPACE to Start",
            "Press ESC to Quit"
        ]

        y = 250
        for instruction in instructions:
            text = self.font_small.render(instruction, True, DARK_GRAY)
            text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, y))
            self.screen.blit(text, text_rect)
            y += 30

    def draw_game(self):
        """绘制游戏界面"""
        self.screen.fill(LIGHT_GRAY)

        # 绘制拼图区域背景
        puzzle_area = pygame.Rect(MARGIN - 2, MARGIN - 2, PUZZLE_SIZE + 4, PUZZLE_SIZE + 4)
        pygame.draw.rect(self.screen, WHITE, puzzle_area)
        pygame.draw.rect(self.screen, BLACK, puzzle_area, 3)

        # 绘制网格线
        for i in range(GRID_SIZE + 1):
            # 垂直线
            x = MARGIN + i * PIECE_SIZE
            pygame.draw.line(self.screen, GRAY, (x, MARGIN), (x, MARGIN + PUZZLE_SIZE), 1)
            # 水平线
            y = MARGIN + i * PIECE_SIZE
            pygame.draw.line(self.screen, GRAY, (MARGIN, y), (MARGIN + PUZZLE_SIZE, y), 1)

        # 绘制拼图块
        for piece in self.pieces:
            if not piece.is_dragging:
                piece.draw(self.screen)

        # 最后绘制正在拖拽的块（确保在最上层）
        if self.dragging_piece:
            self.dragging_piece.draw(self.screen)

        # 绘制统计信息
        self.draw_stats()

        # 绘制预览图
        self.draw_preview()

    def draw_stats(self):
        """绘制游戏统计信息"""
        # 移动次数
        moves_text = self.font_medium.render(f"Moves: {self.moves}", True, BLACK)
        self.screen.blit(moves_text, (SCREEN_WIDTH - 200, 50))

        # 游戏时间
        if self.game_started:
            minutes = self.time_elapsed // 60
            seconds = self.time_elapsed % 60
            time_text = self.font_medium.render(f"Time: {minutes:02d}:{seconds:02d}", True, BLACK)
            self.screen.blit(time_text, (SCREEN_WIDTH - 200, 90))

        # 完成进度
        correct_pieces = sum(1 for piece in self.pieces if piece.is_correct)
        progress_text = self.font_small.render(f"Progress: {correct_pieces}/{len(self.pieces)}", True, BLUE)
        self.screen.blit(progress_text, (SCREEN_WIDTH - 200, 130))

        # 操作提示
        hint_text = self.font_small.render("Press R - Reset", True, GRAY)
        self.screen.blit(hint_text, (SCREEN_WIDTH - 200, 170))

        hint_text2 = self.font_small.render("Press ESC - Menu", True, GRAY)
        self.screen.blit(hint_text2, (SCREEN_WIDTH - 200, 200))

    def draw_preview(self):
        """绘制预览图"""
        preview_size = 150
        preview_x = SCREEN_WIDTH - preview_size - 50
        preview_y = 250

        # 绘制预览图背景
        preview_rect = pygame.Rect(preview_x, preview_y, preview_size, preview_size)
        pygame.draw.rect(self.screen, WHITE, preview_rect)
        pygame.draw.rect(self.screen, BLACK, preview_rect, 2)

        # 缩放并绘制原图
        scaled_image = pygame.transform.scale(self.original_image, (preview_size, preview_size))
        self.screen.blit(scaled_image, preview_rect)

        # 预览标签
        preview_label = self.font_small.render("Preview", True, BLACK)
        label_rect = preview_label.get_rect(center=(preview_x + preview_size // 2, preview_y - 15))
        self.screen.blit(preview_label, label_rect)

    def draw_win_screen(self):
        """绘制胜利界面"""
        self.draw_game()

        # 半透明覆盖层
        overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
        overlay.set_alpha(200)
        overlay.fill(WHITE)
        self.screen.blit(overlay, (0, 0))

        # 胜利文字
        win_text = self.font_large.render("PUZZLE COMPLETED!", True, GREEN)
        win_rect = win_text.get_rect(center=(SCREEN_WIDTH // 2, 200))
        self.screen.blit(win_text, win_rect)

        # 统计信息
        stats = [
            f"Total Moves: {self.moves}",
            f"Time: {self.time_elapsed // 60:02d}:{self.time_elapsed % 60:02d}",
            "",
            "Press SPACE to Play Again",
            "Press ESC for Menu"
        ]

        y = 280
        for stat in stats:
            text = self.font_medium.render(stat, True, BLACK)
            text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, y))
            self.screen.blit(text, text_rect)
            y += 40

    def handle_mouse_down(self, pos):
        """处理鼠标按下事件"""
        if self.state == GameState.PLAYING:
            piece = self.get_piece_at_pos(pos)
            if piece:
                self.dragging_piece = piece
                piece.is_dragging = True
                self.drag_start_pos = pos
                piece.drag_offset = (pos[0] - piece.rect.x, pos[1] - piece.rect.y)

    def handle_mouse_up(self, pos):
        """处理鼠标释放事件"""
        if self.state == GameState.PLAYING and self.dragging_piece:
            # 找到释放位置最近的拼图块
            target_piece = self.get_piece_at_pos(pos)

            if target_piece and target_piece != self.dragging_piece:
                # 交换两个拼图块
                self.swap_pieces(self.dragging_piece, target_piece)
            else:
                # 将拼图块吸附到最近的网格位置
                grid_col = round((pos[0] - MARGIN) / PIECE_SIZE - 0.5)
                grid_row = round((pos[1] - MARGIN) / PIECE_SIZE - 0.5)

                grid_col = max(0, min(GRID_SIZE - 1, grid_col))
                grid_row = max(0, min(GRID_SIZE - 1, grid_row))

                # 检查目标位置是否被占用
                occupied = any(piece.current_pos == (grid_row, grid_col) and piece != self.dragging_piece
                             for piece in self.pieces)

                if not occupied:
                    self.dragging_piece.update_position((grid_row, grid_col))
                    self.moves += 1

            # 重置拖拽状态
            self.dragging_piece.is_dragging = False
            self.dragging_piece = None
            self.drag_start_pos = None

            # 检查是否获胜
            if self.check_win():
                self.state = GameState.WON

    def handle_mouse_motion(self, pos):
        """处理鼠标移动事件"""
        if self.state == GameState.PLAYING and self.dragging_piece:
            # 更新拖拽块的位置
            self.dragging_piece.rect.x = pos[0] - self.dragging_piece.drag_offset[0]
            self.dragging_piece.rect.y = pos[1] - self.dragging_piece.drag_offset[1]

    def reset_game(self):
        """重置游戏"""
        self.create_puzzle_pieces()
        self.shuffle_pieces()
        self.moves = 0
        self.time_elapsed = 0
        self.game_started = True
        self.dragging_piece = None
        self.drag_start_pos = None

    def handle_events(self):
        """处理事件"""
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False

            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    if self.state == GameState.PLAYING or self.state == GameState.WON:
                        self.state = GameState.MENU
                        self.game_started = False
                    else:
                        return False

                elif event.key == pygame.K_SPACE:
                    if self.state == GameState.MENU:
                        self.state = GameState.PLAYING
                        self.reset_game()
                    elif self.state == GameState.WON:
                        self.state = GameState.PLAYING
                        self.reset_game()

                elif event.key == pygame.K_r and self.state == GameState.PLAYING:
                    self.reset_game()

            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:  # 左键
                    self.handle_mouse_down(event.pos)

            elif event.type == pygame.MOUSEBUTTONUP:
                if event.button == 1:  # 左键
                    self.handle_mouse_up(event.pos)

            elif event.type == pygame.MOUSEMOTION:
                self.handle_mouse_motion(event.pos)

        return True

    def update(self):
        """更新游戏状态"""
        if self.state == GameState.PLAYING and self.game_started:
            self.time_elapsed += 1 // FPS  # 假设FPS为60

    def draw(self):
        """绘制游戏"""
        if self.state == GameState.MENU:
            self.draw_menu()
        elif self.state == GameState.PLAYING:
            self.draw_game()
        elif self.state == GameState.WON:
            self.draw_win_screen()

        pygame.display.flip()

    def run(self):
        """运行游戏"""
        running = True

        while running:
            running = self.handle_events()
            self.update()
            self.draw()
            self.clock.tick(FPS)

        pygame.quit()
        sys.exit()

# 运行游戏
if __name__ == "__main__":
    game = PuzzleGame()
    game.run()